#!/bin/bash
# pgBackRest 数据库恢复脚本

# 脚本目录路径
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
COMMON_DIR="${SCRIPT_DIR}/../common"

# 加载公共工具函数
source "${COMMON_DIR}/utils.sh"

# 显示使用说明
show_usage() {
    echo "用法: $0 [选项]"
    echo "选项:"
    echo "  --type=TYPE         恢复类型: full(默认), incr(增量), delta(差异)"
    echo "  --target=TARGET     恢复目标: latest(最新备份, 默认), time(时间点), backup-set(指定备份集)"
    echo "  --set=SET_NAME      指定备份集名称 (当 --target=backup-set 时必需)"
    echo "  --time=TIME_STAMP   指定恢复时间点 (当 --target=time 时必需, 格式: \"YYYY-MM-DD HH:MM:SS\")"
    echo "  --force             强制恢复，不进行确认"
    echo "  --dry-run           仅显示恢复命令，不实际执行"
    echo "  --help              显示此帮助信息"
    echo
    echo "示例:"
    echo "  $0 --type=full                                  # 执行完全恢复到最新备份"
    echo "  $0 --type=delta --target=backup-set --set=20230101-120000F  # 增量恢复到指定备份集"
    echo "  $0 --type=delta --target=time --time=\"2023-01-01 12:00:00\" # 增量恢复到指定时间点"
    echo "  $0 --dry-run --type=full                        # 仅显示完全恢复命令"
}

# 初始化环境
LOG_FILE_PREFIX="pgbackrest_restore"
init
start_timer

# 解析命令行参数
RESTORE_TYPE="full"
RESTORE_TARGET="latest"
BACKUP_SET=""
RESTORE_TIME=""
FORCE=false
DRY_RUN=false

for i in "$@"; do
    case $i in
        --type=*)
            RESTORE_TYPE="${i#*=}"
            ;;
        --target=*)
            RESTORE_TARGET="${i#*=}"
            ;;
        --set=*)
            BACKUP_SET="${i#*=}"
            ;;
        --time=*)
            RESTORE_TIME="${i#*=}"
            ;;
        --force)
            FORCE=true
            ;;
        --dry-run)
            DRY_RUN=true
            ;;
        --help)
            show_usage
            exit 0
            ;;
        *)
            log_error "未知参数: $i"
            show_usage
            exit 1
            ;;
    esac
done

# 验证参数
if [ "$RESTORE_TARGET" = "backup-set" ] && [ -z "$BACKUP_SET" ]; then
    log_error "目标为指定备份集时，必须提供备份集名称 (--set=SET_NAME)"
    show_usage
    exit 1
fi

if [ "$RESTORE_TARGET" = "time" ] && [ -z "$RESTORE_TIME" ]; then
    log_error "目标为时间点还原时，必须提供时间戳 (--time=TIME_STAMP)"
    show_usage
    exit 1
fi

# 检查 PostgreSQL 是否在运行
PG_RUNNING=false
if ps -ef | grep -v grep | grep -q postgres; then
    PG_RUNNING=true
    log_warning "PostgreSQL 服务正在运行，需要先停止服务才能进行恢复"

    if [ "$FORCE" != "true" ] && [ "$DRY_RUN" != "true" ]; then
        read -p "是否停止 PostgreSQL 服务并继续恢复? (y/n): " CONFIRM
        if [ "$CONFIRM" != "y" ] && [ "$CONFIRM" != "Y" ]; then
            log_info "用户取消操作，退出恢复过程"
            exit 0
        fi
    fi

    if [ "$DRY_RUN" != "true" ]; then
        log_info "停止 PostgreSQL 服务..."
        if ! systemctl stop postgresql; then
            log_error "无法停止 PostgreSQL 服务，恢复操作取消"
            exit 1
        fi
        log_success "PostgreSQL 服务已停止"
    else
        log_info "模拟操作: 停止 PostgreSQL 服务"
    fi
fi

# 构建恢复命令
RESTORE_CMD="pgbackrest --stanza=${PGBACKREST_STANZA}"

# 添加恢复类型
if [ "$RESTORE_TYPE" = "delta" ]; then
    RESTORE_CMD="$RESTORE_CMD --delta"
elif [ "$RESTORE_TYPE" = "incr" ]; then
    RESTORE_CMD="$RESTORE_CMD --delta"  # pgBackRest 没有 --incr 选项，使用 --delta
fi

# 添加恢复目标
if [ "$RESTORE_TARGET" = "backup-set" ]; then
    RESTORE_CMD="$RESTORE_CMD --set=${BACKUP_SET}"
elif [ "$RESTORE_TARGET" = "time" ]; then
    RESTORE_CMD="$RESTORE_CMD --type=time --target=\"${RESTORE_TIME}\""
fi

# 完成命令构建
RESTORE_CMD="$RESTORE_CMD restore"

# 显示将要执行的命令
log_info "准备执行以下恢复命令:"
echo "$RESTORE_CMD"

# 确认恢复操作
if [ "$FORCE" != "true" ] && [ "$DRY_RUN" != "true" ]; then
    log_warning "恢复操作将覆盖数据库中的数据，请确保您了解此操作的影响"
    read -p "是否确认继续执行恢复操作? (y/n): " CONFIRM
    if [ "$CONFIRM" != "y" ] && [ "$CONFIRM" != "Y" ]; then
        log_info "用户取消操作，退出恢复过程"
        exit 0
    fi
fi

# 执行恢复
if [ "$DRY_RUN" = "true" ]; then
    log_info "模拟操作: 执行恢复命令"
    log_info "实际恢复将执行以下命令: $RESTORE_CMD"
else
    log_info "开始执行数据库恢复..."

    # 实际执行恢复命令
    if eval "$RESTORE_CMD"; then
        log_success "数据库恢复成功完成"
    else
        log_error "数据库恢复失败"
        exit 1
    fi

    # 如果之前 PostgreSQL 正在运行，则启动服务
    if [ "$PG_RUNNING" = "true" ]; then
        log_info "启动 PostgreSQL 服务..."
        if systemctl start postgresql; then
            log_success "PostgreSQL 服务已启动"
        else
            log_error "无法启动 PostgreSQL 服务，请手动检查服务状态"
            exit 1
        fi
    fi
fi

# 显示执行时间
end_timer

log_info "恢复脚本执行完成"
exit 0
